package com.gitee.ixtf.rsocket.broker

import cn.hutool.log.Log
import com.gitee.ixtf.rsocket.DISCONNECTED_ERROR
import com.gitee.ixtf.rsocket.service.BrokerServiceInstance
import com.gitee.ixtf.rsocket.service.GetBrokerServiceRequest
import com.gitee.ixtf.rsocket.service.GetBrokerServiceResponse
import com.google.common.collect.Maps
import com.google.inject.Singleton
import io.cloudevents.CloudEvent
import io.rsocket.Payload
import io.rsocket.RSocket
import io.rsocket.util.RSocketProxy
import java.security.Principal
import org.reactivestreams.Publisher
import reactor.core.publisher.Flux
import reactor.core.publisher.Mono
import reactor.core.scheduler.Schedulers

private val SERVICE_MAP by lazy { Maps.newConcurrentMap<String, ServiceRSocket>() }

@Singleton
class ServiceManager {
  fun add(principal: Principal, service: String, rsocket: RSocket, cloudEvent: CloudEvent) {
    synchronized(SERVICE_MAP) {
      val serviceRSocket = SERVICE_MAP.compute(service) { k, v -> v ?: ServiceRSocket(k) }
      serviceRSocket!!.add(principal = principal, rsocket = rsocket, cloudEvent = cloudEvent)
    }
  }

  fun pick(split: List<String>): RSocket? {
    val serviceRSocket = SERVICE_MAP[split[0]] ?: return null
    return serviceRSocket.pick(split)
  }

  fun listResponse() = SERVICE_MAP.values.map { it.response() }

  fun getResponse(request: GetBrokerServiceRequest) =
      SERVICE_MAP[request.service]?.response()
          ?: GetBrokerServiceResponse(service = request.service, instances = emptyList())
}

private class ServiceRSocket(private val service: String) : RSocket {
  private val log by lazy { Log.get(this::class.java) }
  private val instanceMap by lazy { Maps.newConcurrentMap<String, InstanceRSocket>() }

  fun response() =
      GetBrokerServiceResponse(
          service = service,
          instances = instanceMap.values.map { it.response() },
      )

  fun pick(split: List<String>? = null): RSocket? =
      when (split?.size) {
        1 -> if (instanceMap.isEmpty()) null else this
        2 -> instanceMap[split[1]]
        else -> instanceMap.values.firstOrNull()
      }

  fun add(principal: Principal, rsocket: RSocket, cloudEvent: CloudEvent) {
    val instance = InstanceRSocket(id = cloudEvent.id, principal = principal, rsocket = rsocket)
    instanceMap.put(instance.id, instance)?.dispose()
    log.debug("$service[${instance.id}] connected by ${principal.name}")
    instance
        .onClose()
        .publishOn(Schedulers.boundedElastic())
        .doAfterTerminate {
          synchronized(SERVICE_MAP) {
            instanceMap.remove(instance.id)
            log.debug("$service[${instance.id}] disconnected")
            if (instanceMap.size == 0) {
              log.debug("$service disconnected")
            }
          }
        }
        .subscribe()
  }

  override fun metadataPush(payload: Payload): Mono<Void> =
      pick()?.metadataPush(payload) ?: Mono.error(DISCONNECTED_ERROR)

  override fun requestResponse(payload: Payload): Mono<Payload> =
      pick()?.requestResponse(payload) ?: Mono.error(DISCONNECTED_ERROR)

  override fun requestStream(payload: Payload): Flux<Payload> =
      pick()?.requestStream(payload) ?: Flux.error(DISCONNECTED_ERROR)

  override fun fireAndForget(payload: Payload): Mono<Void> =
      pick()?.fireAndForget(payload) ?: Mono.error(DISCONNECTED_ERROR)

  override fun requestChannel(payloads: Publisher<Payload>): Flux<Payload> =
      pick()?.requestChannel(payloads) ?: Flux.error(DISCONNECTED_ERROR)
}

private class InstanceRSocket(val id: String, val principal: Principal, rsocket: RSocket) :
    RSocketProxy(rsocket) {
  fun response() = BrokerServiceInstance(id = id)
}
